import rclpy
from rclpy.node import Node
from sensor_msgs.msg import NavSatFix
from geometry_msgs.msg import TwistStamped
from sls_msgs.msg import LlaVelocity
import math
from std_msgs.msg import Float64, Float64MultiArray
from rclpy.qos import QoSProfile, ReliabilityPolicy
from nav_msgs.msg import Odometry
from scipy.spatial.transform import Rotation as R
from functools import partial

class RelayNode(Node):
    def __init__(self):
        super().__init__('lla_velocity_relay')
        self.declare_parameter('ugv_ids', ['ugv1'])
        self.ugv_ids = self.get_parameter('ugv_ids').value
        self.altitude = 0.0
        # 设置 QoS 为 BEST_EFFORT 以兼容传感器数据
        qos = QoSProfile(
            reliability=ReliabilityPolicy.BEST_EFFORT,
            depth=10
        )
        self.ugv_states = {}
        self.sensor_msg_pub = {}
        wheel_radius = 0.225  # 可作为参数，这里硬编码

        for ugv_id in self.ugv_ids:
            # 初始化状态
            self.ugv_states[ugv_id] = {
                'latitude': 0.0,
                'longitude': 0.0,
                'altitude': 0.0,
                'speed_rl': 0.0,
                'speed_rr': 0.0,
                'steer0': 0.0,
                'steer1': 0.0,
                'yaw': None,
                'wheel_radius': wheel_radius
            }
            pub_topic = f'/{ugv_id}/sensors/rtk/out/lla'
            self.create_subscription(
                Float64MultiArray,
                f'/{ugv_id}/actuators/chassis/out/motor_data',
                partial(self.motor_data_callback, ugv_id=ugv_id),
                10
            )
            self.create_subscription(
                Float64MultiArray,
                f'/{ugv_id}/actuators/chassis/out/steer_data',
                partial(self.steering_data_callback, ugv_id=ugv_id),
                10
            )
            self.create_subscription(
                NavSatFix,
                f'/{ugv_id}/gps',
                partial(self.gps_callback, ugv_id=ugv_id),
                10
            )
            self.create_subscription(
                Odometry,
                f'/{ugv_id}/out/odom',
                partial(self.odom_callback, ugv_id=ugv_id),
                10
            )
        

            self.sensor_msg_pub[ugv_id] = self.create_publisher(
                LlaVelocity,
                pub_topic,
                10
            )
        for ugv_id in self.ugv_ids:
            self.create_timer(0.5, partial(self.publish_combined_msg, ugv_id=ugv_id))

    def gps_callback(self, msg:NavSatFix, ugv_id: str):
        state = self.ugv_states[ugv_id]
        state['latitude'] = msg.latitude
        state['longitude'] = msg.longitude
        state['altitude'] = msg.altitude

    def motor_data_callback(self, msg:Float64MultiArray, ugv_id: str):
        state = self.ugv_states[ugv_id]
        state['speed_rl'] = msg.data[0]
        state['speed_rr'] = msg.data[1]
    
    def steering_data_callback(self, msg:Float64MultiArray, ugv_id: str):
        state = self.ugv_states[ugv_id]
        state['steer0'] = msg.data[0]
        state['steer1'] = msg.data[1]
        # self.get_logger().info(f"0{self.steer0},1{self.steer1}")

    def odom_callback(self, msg:Odometry, ugv_id: str):
        state = self.ugv_states[ugv_id]
        q = msg.pose.pose.orientation
    # 四元数格式：[x, y, z, w]
        quat = [q.x, q.y, q.z, q.w]
        yaw = math.acos(q.w)*2
        if q.z < 0:
            yaw = -yaw
    # # 创建 Rotation 对象并转换为欧拉角（'xyz' 或 'rpy' 顺序）
    #     r = R.from_quat(quat)
    #     roll, pitch, yaw = r.as_euler('zyx', degrees=False)  # 输出为弧度
    #     self.get_logger().info(f"[{ugv_id}] RAW QUAT: x={q.x:.4f}, y={q.y:.4f}, z={q.z:.4f}, w={q.w:.4f}")
        self.ugv_states[ugv_id]['yaw'] = yaw

    def publish_combined_msg(self, ugv_id: str):
        state = self.ugv_states[ugv_id]
        # self.get_logger().info(f"steer0={self.steer0}, steer1={self.steer1}")
        msg = LlaVelocity()
        msg.latitude = state['latitude']
        msg.longitude = state['longitude']
        msg.altitude = state['altitude']
        wheel_radius = state['wheel_radius']  # m
        current_speed = wheel_radius * 0.5 * (state['speed_rl'] + state['speed_rr'])

        msg.vel = current_speed
        try:
            msg.yaw = math.degrees(float(state['yaw']))
        except:
            msg.yaw = 0.0
        try:
            msg.vel_heading = math.degrees(msg.yaw)
        except:
            msg.vel_heading = 0.0
        # except:
        #     msg.speed = 0.0
        self.sensor_msg_pub[ugv_id].publish(msg)
    

def main(args=None):
    rclpy.init(args=args)
    node = RelayNode()
    rclpy.spin(node)
    node.destroy_node()
    rclpy.shutdown()
if __name__ == '__main__':
    main()